import streamlit as st
from openai import OpenAI
import pickle
import os
import time
import json
from streamlit import session_state 
os.environ['http_proxy'] = 'http://127.0.0.1:7890'
os.environ["https_proxy"] = "http://127.0.0.1:7890"
#streamlit run dataset.py --server.port 2323
st.set_page_config(
    page_title='问答数据集生成器',
    layout="wide",
    page_icon='😅',
    initial_sidebar_state="expanded",#“auto”或“expanded”或“collapsed”
         menu_items={
         'Get Help': None,
         'Report a bug': None,
         'About': None
     }
)
# 加载问题库
def load_questions(file_path):
    """根据文件中的空行加载分组的问题"""
    if not os.path.exists(file_path):
        st.error(f"文件 {file_path} 不存在")
        create_example_questions()  # 如果文件不存在，创建示例文件
        
    try:
        with open(file_path, "r", encoding='utf-8') as file:
            content = file.read()
        
        # 按空行分割成不同的组
        groups = [g.strip() for g in content.split('\n\n') if g.strip()]
        
        # 处理每个组中的问题
        grouped_questions = []
        for group in groups:
            # 获取这个组中的所有非空问题
            questions = [q.strip() for q in group.split('\n') if q.strip()]
            if questions:  # 如果这个组有问题
                grouped_questions.extend(questions)
        
        return grouped_questions
    except Exception as e:
        st.error(f"加载问题时出错: {str(e)}")
        return []

# 保存问题库
def save_questions(file_path, questions):
    """保存问题时保持原有的分组格式"""
    # 首先读取原有的内容
    original_content = ""
    if os.path.exists(file_path):
        with open(file_path, "r", encoding='utf-8') as file:
            original_content = file.read()
    
    # 保持原有内容不变
    with open(file_path, "w", encoding='utf-8') as file:
        file.write(original_content)

# 使用 GPT3.5 的 API 生成回答
def generate_answer(prompt):
    answer='114514'
    return answer

def save_answers(temp_answers, just_read=False):
    """保存或读取答案"""
    if just_read:
        if os.path.exists("data.pkl"):
            with open("data.pkl", "rb") as file:
                try:
                    answers = pickle.load(file)
                    st.session_state.all_answers = answers
                except Exception as e:
                    st.error(f"读取数据失败：{str(e)}")
                    st.session_state.all_answers = {}
        else:
            st.session_state.all_answers = {}
        return True
    else:
        if os.path.exists("lock"):
            return False
        
        try:
            with open("lock", "w") as lock_file:
                lock_file.write("")

            # 读取现有数据并更新
            if os.path.exists("data.pkl"):
                with open("data.pkl", "rb") as file:
                    try:
                        existing_answers = pickle.load(file)
                    except:
                        existing_answers = {}
            else:
                existing_answers = {}

            # 更新答案
            existing_answers.update(temp_answers)
            
            # 保存更新后的数据
            with open("data.pkl", "wb") as file:
                pickle.dump(existing_answers, file)

            # 更新 session_state
            st.session_state.all_answers = existing_answers

            if os.path.exists("lock"):
                os.remove("lock")
            return True

        except Exception as e:
            st.error(f"保存数据失败：{str(e)}")
            if os.path.exists("lock"):
                os.remove("lock")
            return False

def save_answers_as_json(answers, file_path):
    """保存答案为sessions格式的JSON文件"""
    output = {
        "sessions": []
    }
    
    try:
        # 读取问题文件以获取分组信息
        with open(st.session_state.question_txt, "r", encoding='utf-8') as file:
            content = file.read()
        
        # 按空行分割获取分组
        groups = [g.strip() for g in content.split('\n\n') if g.strip()]
        
        session_id = 0
        for group in groups:
            questions_in_group = [q.strip() for q in group.split('\n') if q.strip()]
            if not questions_in_group:
                continue
                
            session = {
                "sessionID": session_id,
                "chatTurns": {}
            }
            
            has_answers = False
            
            # 为这个组中的每个问题添加答案
            for turn_id, question in enumerate(questions_in_group):
                if question in answers:
                    session["chatTurns"][str(turn_id)] = {
                        "question": question,
                        "answer": answers[question]
                    }
                    has_answers = True
            
            if has_answers:
                output["sessions"].append(session)
                session_id += 1
        
        # 写入文件
        with open(file_path, "w", encoding="utf-8") as file:
            json.dump(output, file, ensure_ascii=False, indent=4)
        
        return True
        
    except Exception as e:
        st.error(f"处理文件时出错: {str(e)}")
        return False

def reset_text_area():
    if session_state.text_area_tittle=="回答：(内容为空则不保存此回答)":
        session_state.text_area_tittle="回答：(内容为空则不保存此回答) "
    elif session_state.text_area_tittle=="回答：(内容为空则不保存此回答) ":
        session_state.text_area_tittle="回答：(内容为空则不保存此回答)"
def get_session_context(question, questions_txt, all_answers):
    """获取当前问题所在session的上下文"""
    try:
        with open(questions_txt, "r", encoding='utf-8') as file:
            content = file.read()
        
        # 按空行分割获取所有session
        sessions = [s.strip() for s in content.split('\n\n') if s.strip()]
        
        # 找到当前问题所在的session
        current_session = None
        for session in sessions:
            questions = [q.strip() for q in session.split('\n') if q.strip()]
            if question in questions:
                current_session = questions
                break
        
        if not current_session:
            return []
        
        # 只获取当前session中，当前问题之前的问答对
        context = []
        for q in current_session:
            if q == question:  # 到达当前问题就停止
                break
            if q in all_answers:  # 只添加已经回答过的问题
                context.append({
                    "question": q,
                    "answer": all_answers[q]
                })
        return context
        
    except Exception as e:
        st.error(f"获取上下文失败: {str(e)}")
    return []

def format_context(session_questions, current_question, answers):
    """格式化上下文信息"""
    context = []
    for q in session_questions:
        if q == current_question:
            break
        if q in answers:
            context.append(f"问：{q}\n答：{answers[q]}")
    
    return "\n\n".join(context)

# 在程序启动时创建示例问题文件（如果不存在）
def create_example_questions():
    if not os.path.exists("questions.txt"):
        example_questions = """为什么春节要吃月饼？
春节吃月饼有什么寓意？
���什么春节吃月饼不可以吃太阳饼？
春节吃月饼是古代遗留下来的吗？

王一博在《》中扮演的角色是什么？
他在《八角笼中》中的表演受到了哪些评价？
表演有哪些亮点？
戏份占比是多少？

陶渊明的《菊花源记》讲了什么？
陶渊明的《菊花源记》的创作背景是什么？
《菊花源记》是陶渊明的哪一篇作品？
陶渊明的《菊花源记》中的主要人物有哪些？"""
        
        with open("questions.txt", "w", encoding="utf-8") as f:
            f.write(example_questions)

def main():
    # 确保questions.txt文件存在
    create_example_questions()
    
    st.title("该死数据集生成器")
   
    PROMPT=st.sidebar.text_input("提示词", value= "请给出以下问题的答案")
    if 'temp_answers' not in session_state:
        session_state.temp_answers={}
    if 'all_answers' not in session_state:
        save_answers(session_state.temp_answers,just_read=True)
        session_state.question_txt="questions.txt"
        session_state.answers_json="answers.json"
        session_state.generated_answer=""
        session_state.text_area_tittle="回答：(内容为空则不保存此回答)"
        session_state.selected_id=0
    session_state.question_txt=st.sidebar.text_input("存有每一条问题的txt", value=session_state.question_txt)
    session_state.answers_json=st.sidebar.text_input("保存回答的json路径", value=session_state.answers_json)
    if 'questions' not in session_state:
        session_state.questions = load_questions(session_state.question_txt)
    selected_questions={}
    for q in range(len(session_state.questions)):
        selected_questions[session_state.questions[q]]=q 
    selectbox_empty = st.empty()
    selected_question = selectbox_empty.selectbox("请选择一个问题：", session_state.questions,index=session_state.selected_id)
    if selected_question:
        session_state.selected_id = selected_questions[selected_question]
        selected_question = selectbox_empty.selectbox("请选择一个问题： ", session_state.questions,index=session_state.selected_id)
        prompt = PROMPT + selected_question
        st.sidebar.write({'预':prompt})
        user_answer_empty = st.empty()
        user_answer = user_answer_empty.text_area(session_state.text_area_tittle, session_state.generated_answer, height=200)


        col1, col2, col3, col4, col5 = st.columns(5)
        with col1:
            if st.button("保存全部回答(未保存超过10个会自动保存的)") or len(session_state.temp_answers) >= 10:
                session_state.selected_id = 0
                if save_answers(session_state.temp_answers):
                    st.success("全部答已保存。")
                    session_state.temp_answers = {}
                else:
                    st.error("保存失败，请稍后重试。经常出现此问题是因为死锁，请删除data.pkl文件后重试。")
                    time.sleep(5)
                st.rerun()
        with col2:
            if st.button("gpt生成回答"):
                session_state.generated_answer = ''
                try:
                    # 获取当前问题的上下文
                    context = get_session_context(
                        selected_question, 
                        st.session_state.question_txt,
                        {**st.session_state.all_answers, **st.session_state.temp_answers}
                    )
                    
                    # 构建带上下文的提示词
                    if context:
                        context_text = "\n\n".join([
                            f"问：{qa['question']}\n答：{qa['answer']}" 
                            for qa in context
                        ])
                        full_prompt = f"""这是一个连续的对话，（注意：回答要简洁，先说结论，再有20字以内的原因）前面的对话内容是：\n\n{context_text}\n\n基于以上对话的上下文，请回答下面的问题：\n{selected_question}"""
                    else:
                        full_prompt = f"请回答这个问题：{selected_question}"

                    # 添加重试机制
                    max_retries = 3
                    retry_delay = 2
                    for attempt in range(max_retries):
                        try:
                            client = OpenAI(
                                api_key="sk-Umo98NwKInK20LOrMFjtxfpjbE9fYfK06K7SukJgdr1KXPTf",
                                base_url="https://api.moonshot.cn/v1"
                            )
                            
                            completion = client.chat.completions.create(
                                model="moonshot-v1-8k",
                                messages=[
                                    {"role": "system", "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手。你会为用户提供连贯的对话回答，注意保持回答与上下文的一致性。还要就是回答简洁"},
                                    {"role": "user", "content": full_prompt}
                                ],
                                temperature=0.3
                            )
                            session_state.generated_answer = completion.choices[0].message.content
                            break
                        except Exception as e:
                            if "rate_limit" in str(e).lower() and attempt < max_retries - 1:
                                time.sleep(retry_delay * (attempt + 1))
                                continue
                            raise e
                    
                    reset_text_area()
                    user_answer = user_answer_empty.text_area(
                        session_state.text_area_tittle, 
                        session_state.generated_answer, 
                        height=200
                    )
                except Exception as e:
                    st.error("生成失败，请稍后重试")
        with col3:
            if st.button("确认此回答(自动下一个)"):
                session_state.generated_answer = ''
                # 保存当前答案
                if user_answer != '':
                    session_state.temp_answers[selected_question] = user_answer
                elif selected_question in session_state.temp_answers:
                    del session_state.temp_answers[selected_question]
                
                # 移动到下一个问题前，先获取当前组的所有问题
                with open(st.session_state.question_txt, "r", encoding='utf-8') as file:
                    content = file.read()
                groups = [g.strip() for g in content.split('\n\n') if g.strip()]
                
                # 找到当前问题所在的组和位置
                current_group_index = -1
                current_question_index = -1
                for group_idx, group in enumerate(groups):
                    questions = [q.strip() for q in group.split('\n') if q.strip()]
                    if selected_question in questions:
                        current_group_index = group_idx
                        current_question_index = questions.index(selected_question)
                        # 如果不是组内最后一个问题，移动到下一个
                        if current_question_index < len(questions) - 1:
                            session_state.selected_id += 1
                        # 如果是组内最后一个问题，移动到下一组的第一个问题
                        else:
                            next_group_questions = []
                            # 查找下一组的问题
                            if group_idx < len(groups) - 1:
                                next_group = groups[group_idx + 1]
                                next_group_questions = [q.strip() for q in next_group.split('\n') if q.strip()]
                            
                            if next_group_questions:
                                # 找到下一组第一个问题的索引
                                first_question = next_group_questions[0]
                                for i, q in enumerate(session_state.questions):
                                    if q == first_question:
                                        session_state.selected_id = i
                                        break
                            else:
                                # 如果没有下一组，回到第一组第一个问题
                                session_state.selected_id = 0
                        break
                
                # 重置文本区域
                reset_text_area()
                user_answer = user_answer_empty.text_area(session_state.text_area_tittle, height=200)
                st.rerun()
        with col4:
            if st.button("上一个问题"):
                session_state.generated_answer=''
                session_state.selected_id-=1
                if session_state.selected_id<0:
                    session_state.selected_id=len(session_state.questions)-1
                reset_text_area()
                user_answer = user_answer_empty.text_area(session_state.text_area_tittle, height=200)
                st.rerun()
        with col5:
            if st.button("下一个问题"):
                session_state.generated_answer=''
                session_state.selected_id+=1
                if session_state.selected_id>=len(session_state.questions):
                    session_state.selected_id=0
                reset_text_area()
                user_answer = user_answer_empty.text_area(session_state.text_area_tittle, height=200)
                st.rerun()
    if st.sidebar.button("读一下"):
        session_state.selected_id=0
        save_answers(session_state.temp_answers,just_read=True)
        session_state.questions = load_questions(session_state.question_txt)
        st.rerun()
    if st.sidebar.button("导出载入的已保存回答为 JSON"):
        if st.session_state.all_answers:
            if save_answers_as_json(st.session_state.all_answers, st.session_state.answers_json):
                st.success("成功导出JSON文件")
            else:
                st.error("导出过程中出现问题，请检查数据")
        else:
            st.error("没有找到任何已保存的答案")
    st.json({"未保存回答：":session_state.temp_answers,"已保存回答：":session_state.all_answers})

if __name__ == "__main__":
    main()
